import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.font_manager import FontProperties
from matplotlib.colors import LinearSegmentedColormap

# 设置全局字体配置
plt.rcParams['font.size'] = 10
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['text.usetex'] = False

# 创建支持中文的字体属性
chinese_font = FontProperties(fname=mpl.font_manager.findfont('SimHei'))

# 创建图形
fig, ax = plt.subplots(figsize=(8, 6), dpi=100)

# ================== 客户端权重分布热力图 ==================
# 模拟数据：100个客户端的数据质量和梯度相似性
np.random.seed(42)
data_quality = np.random.uniform(0.2, 1.0, 100)  # 数据质量 (0.2-1.0)
gradient_similarity = np.random.uniform(0.4, 0.95, 100)  # 梯度相似性 (0.4-0.95)

# 计算权重：w_k = exp(β * sim_k) × Q_k
beta = 0.5
quality_score = 0.6 * data_quality + 0.4 * gradient_similarity
weights = np.exp(beta * gradient_similarity) * quality_score
weights = weights / weights.max()  # 归一化到0-1

# 创建热力图数据
heatmap, xedges, yedges = np.histogram2d(
    data_quality, 
    gradient_similarity, 
    bins=10,
    weights=weights
)

# 创建灰度颜色映射 - 适配黑白印刷
colors = ["#ffffff", "#d9d9d9", "#969696", "#252525"]
cmap = LinearSegmentedColormap.from_list("custom_gray", colors, N=256)

# 绘制热力图
im = ax.imshow(
    heatmap.T, 
    origin='lower', 
    extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]],
    aspect='auto', 
    cmap=cmap,
    interpolation='nearest'
)

# 添加高质量客户端标注 - 使用圆圈
ax.scatter(
    data_quality[weights > 0.8], 
    gradient_similarity[weights > 0.8], 
    marker='o',  # 圆圈标记
    s=50, 
    facecolors='none', 
    edgecolors='black',  # 黑色边框
    linewidths=1.5,
    label='高质量客户端 (权重>0.8)'
)

# 添加低质量客户端标注 - 使用三角形
ax.scatter(
    data_quality[weights < 0.3], 
    gradient_similarity[weights < 0.3], 
    marker='^',  # 三角形标记
    s=50, 
    facecolors='none', 
    edgecolors='black',  # 黑色边框
    linewidths=1.0,
    label='低质量客户端 (权重<0.3)'
)

# 设置属性
ax.set_xlabel('数据质量 (Q_k)', fontsize=12, fontproperties=chinese_font)
ax.set_ylabel('梯度相似性 (sim_k)', fontsize=12, fontproperties=chinese_font)
ax.set_title('图5：客户端权重分布热力图', fontsize=14, fontweight='bold', fontproperties=chinese_font)
ax.grid(False)
ax.legend(loc='lower left', prop=chinese_font, frameon=True, framealpha=0.9)

# 添加颜色条
cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
cbar.set_label('聚合权重 (w_k)', fontsize=10, fontproperties=chinese_font)

# 添加关键区域标注
ax.text(0.85, 0.92, '高质量区域\n平均权重: 0.18', 
        ha='center', fontproperties=chinese_font, 
        bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="black", alpha=0.8))
ax.text(0.35, 0.52, '低质量区域\n平均权重: 0.03', 
        ha='center', fontproperties=chinese_font, 
        bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="black", alpha=0.8))

# 添加技术标注
fig.text(0.5, 0.01, 
         "实验设置: Non-IID参数α=0.5 | 加权策略: w_k = exp(β·sim_k) × Q_k (β=0.5)", 
         ha="center", fontsize=11, style='italic', fontproperties=chinese_font)

# 添加关键结论标注
ax.text(0.5, 1.05, '高质量客户端权重是低质量客户端的6倍', 
        transform=ax.transAxes, ha='center', fontsize=12, 
        fontproperties=chinese_font, bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="black", alpha=0.8))

# 调整布局并保存
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
plt.savefig('dynamic_weighting_heatmap_bw.png', bbox_inches='tight', dpi=300)
plt.show()
